import numpy as np


# Needle算法对齐
class Needle:
    # 分数矩阵
    score_matrix = [
        [10, -5, 0, -5],
        [-5, 10, -5, 0],
        [0, -5, 10, -5],
        [-5, 0, -5, 10]
    ]
    # 碱基映射关系
    bases = {'A': 0, 'C': 1, 'G': 2, 'T': 3}
    # bases = {'a': 0, 'c': 1, 'g': 2, 't': 3}
    extending_gap = -1
    opening_gap = -4

    def match_score(self, alpha, beta, is_pre_gap):  # is_pre_gap, 前一位是否是gap
        if alpha != '-' and beta != '-':
            score, is_gap = self.score_matrix[self.bases[alpha]][self.bases[beta]], False
        else:
            if is_pre_gap:
                score = self.opening_gap
            else:
                score = self.extending_gap
            is_gap = True
        return score, is_gap  # 返回分数和是否是gap

    def needle(self, seq1, seq2):  # 动态规划算法求解
        m, n = len(seq1), len(seq2)  # length of two sequences
        # print(m)
        # print(n)
        # Generate DP table and traceback path pointer matrix
        score = np.zeros((m + 1, n + 1))

        # Calculate DP table
        for i in range(0, m + 1):
            score[i][0] = self.opening_gap * i
        for j in range(0, n + 1):
            score[0][j] = self.opening_gap * j
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                match = score[i - 1][j - 1] + self.match_score(seq1[i - 1], seq2[j - 1], True)[0]  # 最后计算分数不能这样使用
                delete = score[i - 1][j] + self.opening_gap
                insert = score[i][j - 1] + self.opening_gap  # 均使用-4
                score[i][j] = max(match, delete, insert)

        # 这一步的结果已经很好了
        _score = score[m][n]

        # Traceback and compute the alignment
        align1, align2 = '', ''
        i, j = m, n  # start from the bottom right cell
        while i > 0 and j > 0:  # end  the top or the left edge
            score_current = score[i][j]
            score_diagonal = score[i - 1][j - 1]
            score_up = score[i][j - 1]
            score_left = score[i - 1][j]

            if score_current == score_diagonal + self.match_score(seq1[i - 1], seq2[j - 1], True)[0]:
                align1 += seq1[i - 1]
                align2 += seq2[j - 1]
                i -= 1
                j -= 1
            elif score_current == score_left + self.opening_gap:
                align1 += seq1[i - 1]
                align2 += '-'
                i -= 1
            elif score_current == score_up + self.opening_gap:
                align1 += '-'
                align2 += seq2[j - 1]
                j -= 1

        # Finish tracing up to the top left cell
        while i > 0:
            align1 += seq1[i - 1]
            align2 += '-'
            i -= 1
        while j > 0:
            align1 += '-'
            align2 += seq2[j - 1]
            j -= 1

        # finalize(align1, align2)
        align1 = align1[::-1]  # reverse sequence 1
        align2 = align2[::-1]  # reverse sequence 2
        self.finalize(align1, align2)

        del score
        return align1, align2, _score  # 返回对齐的序列和一个差不多分数

    def finalize(self, align1, align2):  # 精确度更高,但我表示不说话, 用不用看心情

        # calculate identity, score and aligned sequences
        symbol = ''
        score = 0
        identity = 0
        is_pre_gap = False
        for i in range(0, len(align1)):
            _score, is_gap = self.match_score(align1[i], align2[i], is_pre_gap)
            is_pre_gap = is_gap  # 更新
            score += _score
            # if two AAs are the same, then output the letter
            if align1[i] == align2[i]:
                symbol = symbol + align1[i]
                identity = identity + 1

            # if they are not identical and none of them is gap
            elif align1[i] != align2[i] and align1[i] != '-' and align2[i] != '-':
                symbol += ' '

            # if one of them is a gap, output a space
            elif align1[i] == '-' or align2[i] == '-':
                symbol += ' '

        # identity = float(identity) / len(align1) * 100
        # print('Identity =', "%3.3f" % identity, 'percent')
        # print('Score =', score)
        # print(align1)
        # print(symbol)
        # print(align2)
        return score, identity  # 返回分数即可,其余无需返回


if __name__ == "__main__":
    _needle = Needle()
    str1 = 'acgtacgatcgacgatacgtagctgatcaaaaa'
    str2 = 'gctagcatgatcgacgtcaaaaaaaaaaaa'
    _needle.needle(str1, str2)
